import numpy as np
import torch
import torch.utils.data as data
import cv2
import os
import skimage.io as io
import h5py
import random
import torchvision.transforms as transforms
import sys

ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "../"))
sys.path.insert(0, ROOT_DIR)

from utils import train_utils, evaluation_utils

torch.multiprocessing.set_sharing_strategy('file_system')


class Gl3d(data.Dataset):
    def __init__(self, config, mode):
        assert mode == 'train' or mode == 'valid'

        self.config = config
        self.mode = mode
        metadir = os.path.join(config['dataset_path'], 'valid') if mode == 'valid' else os.path.join(config['dataset_path'],
                                                                                                  'train')

        pair_num_list = np.loadtxt(os.path.join(metadir, 'pair_num.txt'), dtype=str)
        self.total_pairs = int(pair_num_list[0, 1])
        self.pair_seq_list, self.accu_pair_num = train_utils.parse_pair_seq(pair_num_list)


    def collate_fn(self, batch):

        data = {}
        dtype = ['im1_ori', 'im2_ori', 'coord1','coord2','F1',\
                'im1','im2','matches']
        for key in dtype:
            data[key] = []
        for sample in batch:
            for key in dtype:
                data[key].append(sample[key])

        for key in ['im1_ori', 'im2_ori', 'coord1','coord2','F1',\
                'im1','im2']:
            data[key] = torch.from_numpy(np.stack(data[key])).float()
        for key in ['matches']:
            data[key] = torch.from_numpy(np.stack(data[key])).int()

        # kpt augmentation with random homography

        return data

    def __getitem__(self, index):
        seq = self.pair_seq_list[index]
        index_within_seq = index - self.accu_pair_num[seq]

        with h5py.File(os.path.join(self.config['dataset_path'], seq, 'info.h5py'), 'r') as data:
            R, t = data['dR'][str(index_within_seq)][()], data['dt'][str(index_within_seq)][()]
            egt = np.reshape(
                np.matmul(np.reshape(evaluation_utils.np_skew_symmetric(t.astype('float64').reshape(1, 3)), (3, 3)),
                          np.reshape(R.astype('float64'), (3, 3))), (3, 3))
            # egt = egt / np.linalg.norm(egt)
            K1, K2 = data['K1'][str(index_within_seq)][()], data['K2'][str(index_within_seq)][()]
            size1, size2 = data['size1'][str(index_within_seq)][()], data['size2'][str(index_within_seq)][()]

            img_path1, img_path2 = data['img_path1'][str(index_within_seq)][()][0].decode(), \
            data['img_path2'][str(index_within_seq)][()][0].decode()
            img_name1, img_name2 = img_path1.split('/')[-1], img_path2.split('/')[-1]
            img_path1, img_path2 = os.path.join(self.config['rawdata_path'], img_path1), os.path.join(
                self.config['rawdata_path'], img_path2)
            fea_path1, fea_path2 = os.path.join(self.config['desc_path'], seq, img_name1 + self.config['desc_suffix']), \
                os.path.join(self.config['desc_path'], seq, img_name2 + self.config['desc_suffix'])
            with h5py.File(fea_path1, 'r') as fea1, h5py.File(fea_path2, 'r') as fea2:
                desc1, kpt1, pscore1 = fea1['descriptors'][()], fea1['keypoints'][()][:, :2], fea1['keypoints'][()][:,
                                                                                              2]
                desc2, kpt2, pscore2 = fea2['descriptors'][()], fea2['keypoints'][()][:, :2], fea2['keypoints'][()][:,
                                                                                              2]
                kpt1, kpt2, desc1, desc2 = kpt1[:self.config['num_kpt']], kpt2[:self.config['num_kpt']], desc1[
                                                                                                   :self.config['num_kpt']], desc2[
                                                                                                                          :self.config['num_kpt']]


            corr = data['corr'][str(index_within_seq)][()]
            incorr1, incorr2 = data['incorr1'][str(index_within_seq)][()], data['incorr2'][str(index_within_seq)][()]

        F1 = np.linalg.inv(K2).T @ egt @ np.linalg.inv(K1)
        F1 /= np.linalg.norm(F1)
        F2 = F1.T

        # permute kpt
        im1 = io.imread(img_path1)
        im2 = io.imread(img_path2)
        im1 = cv2.resize(im1, (1024, 1024))
        im2 = cv2.resize(im2, (1024, 1024))
        im1_tensor = torch.from_numpy(im1/255).permute(2,0,1)
        im2_tensor = torch.from_numpy(im2/255).permute(2,0,1)


        return {'im1_ori': im1, 'im2_ori': im2, 'coord1': kpt1, 'coord2': kpt2, 'F1': F1, 'im1': im1_tensor,
                'im2': im2_tensor,'matches':corr}

    def __len__(self):
        return self.total_pairs

